# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
"""
R* 算法提示词管理模块

本模块负责管理 R* 算法中使用的提示词模板和示例。
R* 是一种基于搜索的推理算法，结合了程序辅助推理（Program-of-Thought）方法。

主要功能:
- 加载和管理提示词模板
- 处理少样本学习示例
- 提供格式化指令和后缀
- 支持随机示例选择

核心组件:
- PROMPT_RSTAR: 主要的提示词管理类
- load_json: JSON 文件加载工具函数
- 提示词格式化和示例管理功能
"""
import json
import random

def load_json(path):
    """
    加载 JSON 文件
    
    读取指定路径的 JSON 文件并返回解析后的数据。
    
    参数:
        path: JSON 文件路径
        
    返回:
        解析后的 JSON 数据
        
    异常:
        FileNotFoundError: 当文件不存在时
        json.JSONDecodeError: 当 JSON 格式错误时
    """
    with open(path, 'r') as f:
        data = json.load(f)
    return data


class PROMPT_RSTAR:
    """
    R* 算法提示词管理类
    
    负责管理 R* 算法中使用的所有提示词组件，包括格式化指令、
    后缀模板和少样本学习示例。支持动态加载和随机选择示例。
    
    属性:
        pot_format_instructions: 程序辅助思维（PoT）格式化指令
        pot_suffix: PoT 后缀模板
        few_examples: 少样本学习示例列表
        num_few_shot: 使用的少样本示例数量
        
    主要方法:
    - load_prompt: 从配置文件加载提示词组件
    - random_examples: 随机选择少样本示例
    """
    
    def __init__(self, config):
        """
        初始化 R* 提示词管理器
        
        根据配置加载提示词模板和示例，设置少样本学习参数。
        
        参数:
            config: 配置对象，包含提示词路径和少样本设置
            
        异常:
            AssertionError: 当请求的少样本数量超过可用示例数量时
        """
        self.pot_format_instructions = None  # PoT 格式化指令
        self.pot_suffix = None              # PoT 后缀模板
        self.few_examples = None            # 少样本示例列表
        self.num_few_shot = config.num_few_shot  # 少样本数量
        
        # 加载提示词组件
        self.load_prompt(config)

        # 验证少样本数量不超过可用示例数量
        assert self.num_few_shot <= len(self.few_examples), f"{self.num_few_shot} should less than few_examples."

    def load_prompt(self, config):
        """
        从配置文件加载提示词组件
        
        读取少样本示例文件和提示词模板文件，初始化所有必要的组件。
        
        参数:
            config: 配置对象，包含文件路径信息
            
        加载内容:
        - few_examples: 从 few_shot_path 加载的示例列表
        - pot_format_instructions: PoT 格式化指令
        - pot_suffix: PoT 后缀模板
        """
        # 加载少样本示例
        self.few_examples = load_json(config.few_shot_path)
        
        # 加载提示词模板
        prompt = load_json(config.prompt_path)
        self.pot_format_instructions = prompt['pot_format_instructions']
        self.pot_suffix = prompt['pot_suffix']

    def random_examples(self):
        """
        随机选择少样本学习示例
        
        根据配置的少样本数量，从可用示例中随机选择指定数量的示例。
        特殊处理：当只需要1个示例时，固定选择第1个示例（索引0）。
        
        返回:
            List: 选中的示例列表
            
        选择策略:
        - 如果 num_few_shot = 1，返回第1个示例（如果存在）
        - 否则随机采样指定数量的示例
        """
        if self.num_few_shot == 1:
            # 特殊情况：只需要1个示例时，选择第1个示例（如果存在）
            if len(self.few_examples) > 0:
                return [self.few_examples[0]]
            else:
                return []
        
        # 随机选择指定数量的示例
        selected_examples = random.sample(
            self.few_examples, 
            min(len(self.few_examples), self.num_few_shot)
        )
        return selected_examples